#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import DrEmpower as dr        # 实物机械臂SDK
import actionlib
import threading
import math
import time

from control_msgs.msg import (
    FollowJointTrajectoryAction,
    FollowJointTrajectoryFeedback,
    FollowJointTrajectoryResult
)
from sensor_msgs.msg import JointState
from trajectory_msgs.msg import JointTrajectoryPoint

# ---------------- 配置区（按需调整） ----------------
joint_names = ["joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "joint_6"]  # ROS关节名
joint_ids   = [1, 2, 3, 4, 5, 6]  # 硬件ID，对应 DrEmpower

joint_state_rate_hz = 50      # /joint_states 发布频率
base_speed = 0.5                # 最小转速 (r/min)，设小一点让计算速度生效
speed_scale = 1.0             # 全局速度放大因子，>1 可以整体加速
ramp_frac = 0.5               # 每段速度最大允许变化比例（50%），用于平滑
accel_param = 2               # 下发时的 param 参数（平滑加减速，可试调）
pos_tolerance_deg = 1.0       # 认为到位的角度容差（deg）
use_velocity_field = True     # 优先使用 trajectory 给出的 velocities
dr_retry_limit = 3            # 读写失败重试次数
initial_wait_seconds = 0.6    # execute_cb 开始前等待 joint_states 时间
# ---------------------------------------------------


class RealArmBridge:
    def __init__(self):
        rospy.init_node("real_arm_bridge_node")

        # 串口互斥锁
        self.lock = threading.Lock()

        # joint_states 缓存
        self.latest_positions = [0.0] * len(joint_ids)  # rad
        self.have_jointstate = threading.Event()

        # 上一次速度缓存（平滑用）
        self.prev_speeds = [base_speed] * len(joint_ids)

        # 发布器
        self.js_pub = rospy.Publisher("/joint_states", JointState, queue_size=10)

        # Action Server
        self.server = actionlib.SimpleActionServer(
            "robot_arm_controller/follow_joint_trajectory",
            FollowJointTrajectoryAction,
            execute_cb=self.execute_cb,
            auto_start=False
        )
        self.server.start()
        rospy.loginfo("✅ FollowJointTrajectory ActionServer 已启动")

        # 启动 joint_state 线程
        t = threading.Thread(target=self.joint_state_loop, daemon=True)
        t.start()

        # 初始 joint_states 发布
        rospy.loginfo("🔹 等待第一次 joint_states 读取...")
        if not self.have_jointstate.wait(timeout=3.0):
            rospy.logwarn("⚠️ 未能在 3s 内读取到初始 joint_states，MoveIt 可能报错")
        else:
            self.publish_joint_states_once()
            rospy.loginfo("🔹 初始 joint_states 已发布")

    # ---------- 发布一次 joint_states ----------
    def publish_joint_states_once(self):
        js = JointState()
        js.header.stamp = rospy.Time.now()
        js.name = joint_names
        js.position = list(self.latest_positions)
        self.js_pub.publish(js)

    # ---------- 安全单关节下发 ----------
    def safe_dr_set_angle(self, id_num, angle, speed=base_speed, param=accel_param, mode=1):
        """
        封装 dr.set_angle，带锁和重试
        angle: deg
        speed: SDK 的 speed 参数（r/min）
        """
        for attempt in range(dr_retry_limit):
            try:
                with self.lock:
                    dr.set_angle(id_num=id_num, angle=angle, speed=speed, param=param, mode=mode)
                return
            except Exception as e:
                rospy.logwarn(f"dr.set_angle({id_num}) 第 {attempt+1}/{dr_retry_limit} 次失败: {e}")
                time.sleep(0.05)
        raise RuntimeError(f"dr.set_angle({id_num}) 连续 {dr_retry_limit} 次失败")

    # ---------- 安全读取角度 ----------
    def safe_dr_get_angle(self, jid):
        for attempt in range(dr_retry_limit):
            try:
                with self.lock:
                    ang = dr.get_angle(id_num=jid)
                return ang
            except Exception as e:
                rospy.logwarn(f"dr.get_angle({jid}) 第 {attempt+1}/{dr_retry_limit} 次失败: {e}")
                time.sleep(0.05)
        raise RuntimeError(f"dr.get_angle({jid}) 连续 {dr_retry_limit} 次失败")

    # ---------- joint_states 循环 ----------
    def joint_state_loop(self):
        rate = rospy.Rate(joint_state_rate_hz)
        while not rospy.is_shutdown():
            positions = []
            try:
                for jid in joint_ids:
                    ang_deg = self.safe_dr_get_angle(jid)
                    positions.append(ang_deg * math.pi / 180.0)  # 转 rad
                self.latest_positions = positions
                js = JointState()
                js.header.stamp = rospy.Time.now()
                js.name = joint_names
                js.position = list(positions)
                self.js_pub.publish(js)
                self.have_jointstate.set()
            except Exception as e:
                rospy.logerr(f"读取关节失败（joint_state_loop）：{e}")
                time.sleep(0.2)
            rate.sleep()

    # ---------- 辅助函数：打印每个关节的计算速度 ----------
    def print_joint_speeds(self, speeds_rpm, idx=None):
        """
        speeds_rpm: list[int] 每个关节的最终速度 r/min
        idx: int, 可选 当前轨迹段索引
        """
        if idx is not None:
            rospy.loginfo(f"🟢 段 {idx}: 关节速度 (r/min) = {speeds_rpm}")
        else:
            rospy.loginfo(f"🟢 当前关节速度 (r/min) = {speeds_rpm}")


    # ------------------- 执行轨迹 -------------------
    def execute_cb(self, goal):
        rospy.loginfo("📥 收到轨迹，共 %d 个点", len(goal.trajectory.points))
        result = FollowJointTrajectoryResult()
        result.error_code = FollowJointTrajectoryResult.SUCCESSFUL

        feedback = FollowJointTrajectoryFeedback()
        feedback.joint_names = joint_names

        rospy.sleep(initial_wait_seconds)
        if not self.have_jointstate.is_set():
            rospy.logwarn("⚠️ 执行时尚未收到 joint_states，可能 MoveIt 认为无效")

        points = goal.trajectory.points
        if len(points) < 2:
            rospy.logwarn("⚠️ 轨迹点不足，直接完成")
            self.server.set_succeeded(result)
            return

        for idx in range(1, len(points)):
            if self.server.is_preempt_requested():
                rospy.loginfo("⏹️ 收到 preempt 请求，停止执行")
                self.server.set_preempted()
                return

            prev_point = points[idx - 1]
            curr_point = points[idx]

            dt = (curr_point.time_from_start - prev_point.time_from_start).to_sec()
            if dt <= 0:
                dt = 0.01

            target_deg = []
            speed_list = []

            # 计算并下发每个关节的速度（r/min）
            for j_idx, jid in enumerate(joint_ids):
                p_prev = prev_point.positions[j_idx]
                p_curr = curr_point.positions[j_idx]

                # 优先使用 velocity 字段
                vel_rad_s = None
                if use_velocity_field and hasattr(curr_point, "velocities") and len(curr_point.velocities) > j_idx:
                    vel = curr_point.velocities[j_idx]
                    if vel is not None and abs(vel) > 1e-6:
                        vel_rad_s = vel
                if vel_rad_s is None:
                    vel_rad_s = (p_curr - p_prev) / dt

                vel_deg_s = abs(vel_rad_s) * 180.0 / math.pi
                rpm = (vel_deg_s / 6.0) * speed_scale

                # 平滑 ramp
                prev_sp = self.prev_speeds[j_idx] if self.prev_speeds else base_speed
                if prev_sp <= 0:
                    prev_sp = base_speed
                min_allowed = max(prev_sp * (1 - ramp_frac), base_speed)
                max_allowed = prev_sp * (1 + ramp_frac)
                rpm_smoothed = min(max(rpm, min_allowed), max_allowed)

                cmd_speed = max(base_speed, rpm_smoothed)

                # 下发到机械臂
                try:
                    self.safe_dr_set_angle(id_num=jid, angle=p_curr * 180.0 / math.pi,
                                        speed=cmd_speed, param=accel_param, mode=1)
                except Exception as e:
                    rospy.logerr(f"❌ dr.set_angle 失败: {e}")
                    result.error_code = FollowJointTrajectoryResult.INVALID_GOAL
                    self.server.set_aborted(result)
                    return

                self.prev_speeds[j_idx] = cmd_speed
                target_deg.append(round(p_curr * 180.0 / math.pi, 2))
                speed_list.append(round(cmd_speed, 2))

            # 输出速度
            self.print_joint_speeds(speed_list, idx=idx)


            rospy.loginfo(f"⬆️ 段 {idx}/{len(points)-1}: targets={target_deg}, dt={dt:.3f}s, speeds(r/min)={speed_list}")

            # 等待该段完成
            t_deadline = rospy.Time.now() + rospy.Duration.from_sec(dt)
            while rospy.Time.now() < t_deadline and not rospy.is_shutdown():
                # 发布反馈
                js = JointState()
                js.header.stamp = rospy.Time.now()
                js.name = joint_names
                js.position = list(self.latest_positions)
                self.js_pub.publish(js)

                feedback.header = js.header
                feedback.actual.positions = js.position
                feedback.desired.positions = curr_point.positions
                feedback.error.positions = [a - b for a, b in zip(js.position, curr_point.positions)]
                self.server.publish_feedback(feedback)

                # 检查是否到位
                all_reached = True
                for j_idx, p_curr in enumerate(curr_point.positions):
                    actual_deg = self.latest_positions[j_idx] * 180.0 / math.pi
                    target_deg_val = p_curr * 180.0 / math.pi
                    if abs(actual_deg - target_deg_val) > pos_tolerance_deg:
                        all_reached = False
                        break
                if all_reached:
                    break

                rospy.sleep(0.02)  # 50Hz 更新

        rospy.loginfo("✅ 轨迹执行完成（逐点下发，平滑处理）")
        self.server.set_succeeded(result)


if __name__ == "__main__":
    try:
        RealArmBridge()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
